load("//bazel:api.bzl", "mojo_filecheck_test", "mojo_test")

_FILECHECK_TESTS = [
    "test_memory.mojo",
    "test_sharedmem_async_cp.mojo",
    "test_tma_ops.mojo",
]

_OTHER_TESTS = [
    "test_shared_mem_barrier.mojo",
    "test_async_cpy_wait_group.mojo",
    "test_buffer_io.mojo",
    "test_copy_async.mojo",
    "test_ldg_intrinsics.mojo",
    "test_ldmatrix_fp8.mojo",
    "test_load_cache.mojo",
    "test_semaphore_reduction.mojo",
    "test_tma.mojo",
    "test_tcgen05.mojo",
    "test_tma_block_reduce.mojo",
]

_EXTRA_CONSTRAINTS = {
    "test_external_and_static_shared_mem.mojo": select({
        "//:apple_gpu": ["@platforms//:incompatible"],
        "//conditions:default": [],
    }),  # FIXME: MOCO-2397
}

mojo_filecheck_test(
    name = "test_shared_mem_barrier.mojo.test",
    srcs = ["test_shared_mem_barrier.mojo"],
    enable_assertions = False,
    tags = ["gpu"],
    target_compatible_with = ["//:has_gpu"],
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

[
    mojo_filecheck_test(
        name = src + ".test",
        srcs = [src],
        tags = ["gpu"] + (
            ["manual"] if src in "test_sharedmem_async_cp.mojo" else []  # TODO: KERN-1531 this hangs with regular builds after cubin compilation, Hangs with debug mode Issue #24921
        ),
        target_compatible_with = ["//:has_gpu"],
        deps = [
            "//max:layout",
            "//max:linalg",
            "@mojo//:stdlib",
        ],
    )
    for src in _FILECHECK_TESTS
]

[
    mojo_test(
        name = src + ".test",
        srcs = [src],
        tags = ["gpu"],
        target_compatible_with = ["//:nvidia_gpu"],
        deps = [
            "//max:layout",
            "//max:linalg",
            "@mojo//:stdlib",
        ],
    )
    for src in [
        "test_async_cpy_wait_group.mojo",
        "test_ldg_intrinsics.mojo",
        "test_copy_async.mojo",
        "test_semaphore_reduction.mojo",  # FIXME: KERN-1377 and move to the globbed mojo_test
        "test_load_cache.mojo",  # FIXME: KERN-1377 and move to the globbed mojo_test
    ]
]

mojo_test(
    name = "test_buffer_io.mojo.test",
    srcs = ["test_buffer_io.mojo"],
    tags = ["gpu"],
    target_compatible_with = ["//:amd_gpu"],
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

mojo_test(
    name = "test_ldmatrix_fp8.mojo.test",
    srcs = ["test_ldmatrix_fp8.mojo"],
    tags = ["gpu"],
    target_compatible_with = ["//:h100_gpu"],
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

mojo_test(
    name = "test_tcgen05.test",
    srcs = ["test_tcgen05.mojo"],
    tags = ["gpu"],
    target_compatible_with = ["//:b200_gpu"],
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

mojo_test(
    name = "test_tma_block_reduce.mojo.test",
    srcs = ["test_tma_block_reduce.mojo"],
    tags = ["gpu"],
    target_compatible_with = select({
        "//:h100_gpu": [],
        "//:b200_gpu": [],
        "//conditions:default": ["@platforms//:incompatible"],
    }),
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

mojo_filecheck_test(
    name = "test_tma.mojo.test",
    srcs = ["test_tma.mojo"],
    tags = ["gpu"],
    target_compatible_with = ["//:h100_gpu"],
    deps = [
        "//max:layout",
        "//max:linalg",
        "@mojo//:stdlib",
    ],
)

[
    mojo_test(
        name = src + ".test",
        srcs = [src],
        tags = ["gpu"],
        target_compatible_with = ["//:has_gpu"] + _EXTRA_CONSTRAINTS.get(src, []),
        deps = [
            "//max:layout",
            "//max:linalg",
            "@mojo//:stdlib",
        ],
    )
    for src in glob(
        ["**/*.mojo"],
        exclude = _FILECHECK_TESTS + _OTHER_TESTS,
    )
]

filegroup(
    name = "test-sources",
    srcs = glob(["**/*.mojo"]),
    visibility = ["//utils/debugging/gpu-build-benchmarking:__subpackages__"],
)
